package com.xzx.flink.streamapi.state;

import com.xzx.flink.bean.ClickEvent;
import com.xzx.flink.streamapi.source.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;

/**
 * 对用户点击事件流每 5 个数据统计一次平均时间戳
 *
 * @version 1.0
 * @auther xinzhixuan
 * @date 2022/5/14 20:27
 */
public class AggregatingState_01_Demo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.addSource(new ClickSource()).assignTimestampsAndWatermarks(WatermarkStrategy.<ClickEvent>forMonotonousTimestamps().withTimestampAssigner(new SerializableTimestampAssigner<ClickEvent>() {
            @Override
            public long extractTimestamp(ClickEvent element, long recordTimestamp) {
                return element.getTimestamp();
            }
        })).keyBy(ClickEvent::getUser).process(new KeyedProcessFunction<String, ClickEvent, String>() {

            private ValueState<Long> countState;
            private AggregatingState<ClickEvent, Long> aggregatingState;

            @Override
            public void open(Configuration parameters) throws Exception {
                countState = getRuntimeContext().getState(new ValueStateDescriptor<>("count", Long.class));
                aggregatingState = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<>("avg", new AggregateFunction<ClickEvent, Tuple2<Long, Long>, Long>() {
                    @Override
                    public Tuple2<Long, Long> createAccumulator() {
                        return Tuple2.of(0L, 0L);
                    }

                    @Override
                    public Tuple2<Long, Long> add(ClickEvent value, Tuple2<Long, Long> accumulator) {
                        return Tuple2.of(value.getTimestamp() + accumulator.f0, 1 + accumulator.f1);
                    }

                    @Override
                    public Long getResult(Tuple2<Long, Long> accumulator) {
                        if (accumulator.f1 == 0) {
                            return 0L;
                        }
                        return accumulator.f0 / accumulator.f1;
                    }

                    @Override
                    public Tuple2<Long, Long> merge(Tuple2<Long, Long> a, Tuple2<Long, Long> b) {
                        return Tuple2.of(a.f0 + b.f0, a.f1 + b.f1);
                    }
                }, Types.TUPLE(Types.LONG, Types.LONG)));
            }

            @Override
            public void processElement(ClickEvent value, KeyedProcessFunction<String, ClickEvent, String>.Context ctx, Collector<String> out) throws Exception {
                aggregatingState.add(value);
                Long count = countState.value();
                if (count == null) {
                    countState.update(1L);
                } else {
                    if (count +1 <5) {
                        countState.update(count + 1L);
                    } else {
                        //达到了5次，输出一次
                        Long avg = aggregatingState.get();
                        out.collect("user:" + value.getUser() + "\t avg ts:" + new Timestamp(avg));
                        countState.clear();
                        aggregatingState.clear();
                    }
                }
            }

        }).print();

        env.execute(AggregatingState_01_Demo.class.getSimpleName());
    }
}
